#!/usr/bin/env python
#-*- coding:utf8 -*-
import torch
from diffusers import FluxPipeline
import os

# 设置输出目录
output_dir = "output"
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

# 直接使用diffusers库而不是modelscope
pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.float16)
# 如果GPU内存不足，可以使用CPU卸载
pipe.enable_model_cpu_offload()

# 生成图像
prompt = "a beautiful sunset over the ocean, photorealistic, high quality"
image = pipe(prompt).images[0]

# 保存图像
output_path = os.path.join(output_dir, "generated_image.png")
image.save(output_path)
print(f"Image generated and saved to: {output_path}")

if __name__ == '__main__':
    prompt = "A cat holding a sign that says hello world"
    image = pipe(
        prompt,
        height=1024,
        width=1024,
        guidance_scale=3.5,
        num_inference_steps=50,
        max_sequence_length=512,
        generator=torch.Generator("cpu").manual_seed(0)
    ).images[0]
    image.save("flux-dev.png")

